CustomOp: Unify aiter impl into GroupedTopk#31221
CustomOp: Unify aiter impl into GroupedTopk#31221ProExpertProg merged 1 commit intovllm-project:mainfrom
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the ROCm AITer implementation for GroupedTopk by encapsulating the platform-specific logic within the GroupedTopk custom operator. This is achieved by introducing a forward_hip method in GroupedTopk and moving the AITer-related checks and function calls there. Consequently, the FusedMoE.select_experts method is simplified, removing conditional logic and direct calls to rocm_aiter_grouped_topk. The changes improve code structure and encapsulation, making the FusedMoE layer more platform-agnostic. The implementation appears correct and is a good improvement.
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
|
@MengqingCao please take a review, thx. |
|
LGTM |
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
6 participants
Purpose
a follow-up of #29575 (comment)
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.